import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import os
from matplotlib import font_manager

# 设置随机种子以确保结果可复现
torch.manual_seed(42)
np.random.seed(42)
# 设置输出路径
output_path = r"D:\pycharm\HomeWork\output"
plots_path = os.path.join(output_path, 'prediction_plots')
# 创建保存图片的文件夹
if not os.path.exists(plots_path):
    os.makedirs(plots_path)
# 设置中文字体
try:
    # 优先使用微软雅黑
    plt.rcParams['font.family'] = ['Microsoft YaHei']
except:
    # 如果没有微软雅黑，则使用系统默认中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
# 读取数据
df = pd.read_excel('./data/GDP.xlsx')

# 准备特征
features = ['GDP（万亿）', 'GDP增长率（%）', '第一产业GDP（万亿元）',
            '第二产业GDP（万亿元）', '第三产业GDP（万亿元）']

# 数据标准化
scaler = MinMaxScaler()
scaled_data = scaler.fit_transform(df[features])


# 准备序列数据
def create_sequences(data, seq_length):
    X, y = [], []
    for i in range(len(data) - seq_length):
        X.append(data[i:(i + seq_length)])
        y.append(data[i + seq_length])
    return torch.FloatTensor(X), torch.FloatTensor(y)


# 定义LSTM模型
class LSTMPredictor(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(LSTMPredictor, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        # LSTM层
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        # 全连接层
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # 初始化隐藏状态和细胞状态
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        # LSTM前向传播
        out, _ = self.lstm(x, (h0, c0))
        # 全连接层
        out = self.fc(out[:, -1, :])
        return out


# 设置参数
seq_length = 5# 使用过去5年的数据预测
input_size = len(features) # 输入特征数
hidden_size = 100# LSTM隐藏层大小
num_layers = 2 # LSTM层数
output_size = len(features) # 输出特征数
num_epochs = 100 # 训练轮数
learning_rate = 0.01# 学习率

# 创建训练数据
X, y = create_sequences(scaled_data, seq_length)

# 初始化模型
model = LSTMPredictor(input_size, hidden_size, num_layers, output_size)
criterion = nn.MSELoss()# 损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)# 优化器

# 训练模型
model.train()
for epoch in range(num_epochs):
    optimizer.zero_grad()# 清空梯度
    outputs = model(X)# 前向传播
    loss = criterion(outputs, y)# 计算损失
    loss.backward() # 反向传播
    optimizer.step() # 更新参数

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 预测未来两年
model.eval()
with torch.no_grad():
    last_sequence = torch.FloatTensor(scaled_data[-seq_length:]).unsqueeze(0)
    predictions = []

    for _ in range(2):  # 预测未来两年
        next_pred = model(last_sequence)
        predictions.append(next_pred.squeeze().numpy())
        last_sequence = torch.cat((last_sequence[:, 1:, :],
                                   next_pred.unsqueeze(1)), dim=1)

# 反标准化预测结果
predictions = scaler.inverse_transform(predictions)

# 创建预测结果DataFrame
future_years = [2024, 2025]
pred_df = pd.DataFrame(predictions, columns=features, index=future_years)

print("\n预测结果：")
print(pred_df)

# 分别保存每个指标的预测图
for i, feature in enumerate(features):
    plt.figure(figsize=(10, 6))

    # 设置网格线
    plt.grid(True, linestyle='--', alpha=0.7)

    # 绘制历史数据和预测数据
    plt.plot(df['年份'], df[feature], 'b-', label='历史数据', marker='o')
    plt.plot(future_years, pred_df[feature], 'r--', label='预测数据', marker='s')

    # 设置标题和标签
    plt.title(feature, fontsize=12, pad=10)
    plt.xlabel('年份', fontsize=10)
    plt.ylabel('数值', fontsize=10)
    plt.legend(fontsize=10)

    # 添加数值标签
    # last_historical = df[feature].iloc[-1]
    # plt.annotate(f'{last_historical:.2f}',
    #              xy=(df['年份'].iloc[-1], last_historical),
    #              xytext=(10, 10), textcoords='offset points',
    #              fontsize=9)

    for year, val in zip(future_years, pred_df[feature]):
        plt.annotate(f'{val:.2f}',
                     xy=(year, val),
                     xytext=(10, -10), textcoords='offset points',
                     fontsize=9)

    # 调整布局
    plt.tight_layout()

    # 保存单个图片
    filename = os.path.join(plots_path, f'{feature}_prediction.png')
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close()

# 创建组合图
fig, axes = plt.subplots(3, 2, figsize=(15, 12))
axes = axes.ravel()

for i, feature in enumerate(features):
    ax = axes[i]
    # 设置网格线
    ax.grid(True, linestyle='--', alpha=0.7)

    # 绘制历史数据和预测数据
    ax.plot(df['年份'], df[feature], 'b-', label='历史数据', marker='o')
    ax.plot(future_years, pred_df[feature], 'r--', label='预测数据', marker='s')

    # 设置标题和标签
    ax.set_title(feature, fontsize=10, pad=10)
    ax.set_xlabel('年份', fontsize=9)
    ax.set_ylabel('数值', fontsize=9)
    ax.legend(fontsize=8)

    # 添加数值标签
    # last_historical = df[feature].iloc[-1]
    # ax.annotate(f'{last_historical:.2f}',
    #             xy=(df['年份'].iloc[-1], last_historical),
    #             xytext=(10, 10), textcoords='offset points',
    #             fontsize=8)

    for year, val in zip(future_years, pred_df[feature]):
        ax.annotate(f'{val:.2f}',
                    xy=(year, val),
                    xytext=(10, -10), textcoords='offset points',
                    fontsize=8)

# 调整子图之间的间距
plt.tight_layout()

# 保存组合图
plt.savefig(os.path.join(plots_path, 'combined_predictions.png'), dpi=300, bbox_inches='tight')
plt.show()

# 保存预测结果到Excel
excel_path = os.path.join(output_path, 'predictions_2024_2025.xlsx')
pred_df.to_excel(excel_path)

# 打印预测结果
print("\n预测结果：")
pd.set_option('display.float_format', lambda x: '%.2f' % x)
print(pred_df)

print(f"\n所有文件已保存至: {output_path}")
print(f"图片保存在: {plots_path}")
print(f"Excel文件保存在: {excel_path}")